/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkPs/MkPsFileSource.h"
#include "MkPs/MkPsDemuxer.h"
#include "MkPs/MkPsMuxer.h"
#include "MkPs/MkPsStreamSource.h"
#include "MkEvent/MkTimer.h"
#include "MkUtil/MkFile.h"
#include <gtest/gtest.h>

extern void WriteBufferListToString(FILE* fp, const CMkBufferList& BufList);
TEST(MkPsUt, PsMuxerTest)
{
    MkString H264FileName = MkString(MkUtSourcePath) + "/test.264";
    MkString H265FileName = MkString(MkUtSourcePath) + "/test.265";
    MkString AacFileName = MkString(MkUtSourcePath) + "/test.aac";
    FILE *fp264 = fopen(H264FileName.c_str(), "rb");
    FILE *fp265 = fopen(H265FileName.c_str(), "rb");
    FILE *fpAac = fopen(AacFileName.c_str(), "rb");
    EXPECT_NE(nullptr, fp264);
    EXPECT_NE(nullptr, fp265);
    EXPECT_NE(nullptr, fpAac);
    CMkH264FileSource H264FileSource(fp264, nullptr);
    CMkH265FileSource H265FileSource(fp265, nullptr);
    CMkAacFileSource AacFileSource(fpAac, nullptr);
    Uint32 ErrorCode = H264FileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(NoneError, ErrorCode);
    ErrorCode = H265FileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(NoneError, ErrorCode);
    ErrorCode = AacFileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(NoneError, ErrorCode);

    MkVector<MkEsPacket> VecPacket264;
    ErrorCode = H264FileSource.GetNextFrame(VecPacket264);
    EXPECT_EQ(NoneError, ErrorCode);

    MkVector<MkEsPacket> VecPacket265;
    ErrorCode = H265FileSource.GetNextFrame(VecPacket265);
    EXPECT_EQ(NoneError, ErrorCode);

    MkVector<MkEsPacket> VecPacketAac;
    ErrorCode = AacFileSource.GetNextFrame(VecPacketAac);
    EXPECT_EQ(NoneError, ErrorCode);

    MkString strModuleName;
    CMkFile::GetModuleFilePath(strModuleName);
    MkString H264AacPsFileName = strModuleName + "/h264aac.ps";
    MkString H264PsFileName = strModuleName + "/h264.ps";
    MkString H265PsFileName = strModuleName + "/h265.ps";
    MkString AacPsFileName = strModuleName + "/aac.ps";
    MkString H265AacPsFileName = strModuleName + "/h265aac.ps";
    FILE *fpH264AacPs = fopen(H264AacPsFileName.c_str(), "wb");
    FILE *fpH264Ps = fopen(H264PsFileName.c_str(), "wb");
    FILE *fpH265Ps = fopen(H265PsFileName.c_str(), "wb");
    //FILE *fpAacPs = fopen(AacPsFileName.c_str(), "wb");
    FILE *fpH265AacPs = fopen(H265AacPsFileName.c_str(), "wb");
    //packet 264 and aac to ps
    {
        Uint8 pBuf[Len1K] = { 0 };
        CMkBufferList BufferList;
        for (int i = 0; i < 100; i++) {
            Uint32 nUsed = 0;
            MkVector<MkEsPacket>::iterator it = VecPacket264.begin();
            for (; it != VecPacket264.end(); it++) {
                ErrorCode = CMkPsMuxer::PacketEsPacket(H264FileSource.GetVideoParam(), AacFileSource.GetAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
                EXPECT_EQ(NoneError, ErrorCode);
                WriteBufferListToString(fpH264AacPs, BufferList);
                BufferList.Clear();
            }

            it = VecPacketAac.begin();
            for (; it != VecPacketAac.end(); it++) {
                ErrorCode = CMkPsMuxer::PacketEsPacket(H264FileSource.GetVideoParam(), AacFileSource.GetAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
                EXPECT_EQ(NoneError, ErrorCode);
                WriteBufferListToString(fpH264AacPs, BufferList);
                BufferList.Clear();
            }
        }
    }

    //packet 265 and aac to ps
    {
        Uint8 pBuf[Len1K] = { 0 };
        CMkBufferList BufferList;
        for (int i = 0; i < 100; i++) {
            Uint32 nUsed = 0;
            MkVector<MkEsPacket>::iterator it = VecPacket265.begin();
            for (; it != VecPacket265.end(); it++) {
                ErrorCode = CMkPsMuxer::PacketEsPacket(H265FileSource.GetVideoParam(), AacFileSource.GetAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
                EXPECT_EQ(NoneError, ErrorCode);
                WriteBufferListToString(fpH265AacPs, BufferList);
                BufferList.Clear();
            }

            it = VecPacketAac.begin();
            for (; it != VecPacketAac.end(); it++) {
                ErrorCode = CMkPsMuxer::PacketEsPacket(H265FileSource.GetVideoParam(), AacFileSource.GetAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
                EXPECT_EQ(NoneError, ErrorCode);
                WriteBufferListToString(fpH265AacPs, BufferList);
                BufferList.Clear();
            }
        }
    }

    //packet 264 to ps
    {
        Uint8 pBuf[Len1K] = { 0 };
        CMkBufferList BufferList;
        for (int i = 0; i < 100; i++) {
            Uint32 nUsed = 0;
            MkVector<MkEsPacket>::iterator it = VecPacket264.begin();
            for (; it != VecPacket264.end(); it++) {
                ErrorCode = CMkPsMuxer::PacketEsPacket(H264FileSource.GetVideoParam(), CMkAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
                EXPECT_EQ(NoneError, ErrorCode);
                WriteBufferListToString(fpH264Ps, BufferList);
                BufferList.Clear();
            }
        }
    }

    //packet 265 to ps
    {
        Uint8 pBuf[Len1K] = { 0 };
        CMkBufferList BufferList;
        for (int i = 0; i < 100; i++) {
            Uint32 nUsed = 0;
            MkVector<MkEsPacket>::iterator it = VecPacket265.begin();
            for (; it != VecPacket265.end(); it++) {
                ErrorCode = CMkPsMuxer::PacketEsPacket(H265FileSource.GetVideoParam(), CMkAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
                EXPECT_EQ(NoneError, ErrorCode);
                WriteBufferListToString(fpH265Ps, BufferList);
                BufferList.Clear();
            }
        }
    }

    //packet aac to ps
    //{
    //    Uint8 pBuf[Len1K] = { 0 };
    //    CMkBufferList BufferList;
    //    for (int i = 0; i < 100; i++) {
    //        Uint32 nUsed = 0;
    //        MkVector<MkEsPacket>::iterator it = VecPacketAac.begin();
    //        for (; it != VecPacketAac.end(); it++) {
    //            ErrorCode = CMkPsMuxer::PacketEsPacket(CMkVideoParam(), AacFileSource.GetAudioParam(), *it, pBuf, Len1K, nUsed, BufferList);
    //            EXPECT_EQ(NoneError, ErrorCode);
    //            WriteBufferListToString(fpAacPs, BufferList);
    //            BufferList.Clear();
    //        }
    //    }
    //}
    fclose(fpH264AacPs);
    fclose(fpH265AacPs);
    fclose(fpH265Ps);
    fclose(fpH264Ps);
    //fclose(fpAacPs);
}

TEST(MkPsUt, PsFileSourceTest)
{
    MkString strModuleName;
    CMkFile::GetModuleFilePath(strModuleName);
    MkString H264AacPsFileName = strModuleName + "/h264aac.ps";
    MkString H264PsFileName = strModuleName + "/h264.ps";
    MkString H265PsFileName = strModuleName + "/h265.ps";
    MkString AacPsFileName = strModuleName + "/aac.ps";
    MkString H265AacPsFileName = strModuleName + "/h265aac.ps";
    //ps file source test
    FILE *fpH264AacPsr = fopen(H264AacPsFileName.c_str(), "rb");
    FILE *fpH264Psr = fopen(H264PsFileName.c_str(), "rb");
    FILE *fpH265Psr = fopen(H265PsFileName.c_str(), "rb");
    //FILE *fpAacPsr = fopen(AacPsFileName.c_str(), "rb");
    FILE *fpH265AacPsr = fopen(H265AacPsFileName.c_str(), "rb");

    CMkPsFileSource PsH264AacFileSource(fpH264AacPsr, nullptr);
    CMkPsFileSource PsH265AacFileSource(fpH265AacPsr, nullptr);
    CMkPsFileSource PsH264FileSource(fpH264Psr, nullptr);
    CMkPsFileSource PsH265FileSource(fpH265Psr, nullptr);
    //CMkPsFileSource PsAacFileSource(fpAacPsr, nullptr);

    PsH264AacFileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(MkCodecH264, PsH264AacFileSource.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecAac, PsH264AacFileSource.GetAudioParam().GetCodec());
    while (1) {
        MkVector<MkEsPacket> VecEsPacket;
        PsH264AacFileSource.GetNextFrame(VecEsPacket);
        if (VecEsPacket.empty()
            || VecEsPacket.begin()->BufferList.Empty()) {
            break;
        }
    }

    PsH265AacFileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(MkCodecH265, PsH265AacFileSource.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecAac, PsH265AacFileSource.GetAudioParam().GetCodec());
    while (1) {
        MkVector<MkEsPacket> VecEsPacket;
        PsH265AacFileSource.GetNextFrame(VecEsPacket);
        if (VecEsPacket.empty()
            || VecEsPacket.begin()->BufferList.Empty()) {
            break;
        }
    }

    PsH264FileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(MkCodecH264, PsH264FileSource.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecUnknown, PsH264FileSource.GetAudioParam().GetCodec());
    while (1) {
        MkVector<MkEsPacket> VecEsPacket;
        PsH264FileSource.GetNextFrame(VecEsPacket);
        if (VecEsPacket.empty()
            || VecEsPacket.begin()->BufferList.Empty()) {
            break;
        }
    }

    PsH265FileSource.StartMediaSource(nullptr, nullptr);
    EXPECT_EQ(MkCodecH265, PsH265FileSource.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecUnknown, PsH265FileSource.GetAudioParam().GetCodec());
    while (1) {
        MkVector<MkEsPacket> VecEsPacket;
        PsH265FileSource.GetNextFrame(VecEsPacket);
        if (VecEsPacket.empty()
            || VecEsPacket.begin()->BufferList.Empty()) {
            break;
        }
    }

    //PsAacFileSource.StartMediaSource(nullptr, nullptr);
    //EXPECT_EQ(MkCodecUnknown, PsAacFileSource.GetVideoParam().GetCodec());
    //EXPECT_EQ(MkCodecAac, PsAacFileSource.GetAudioParam().GetCodec());
    //while (1) {
    //    MkVector<MkEsPacket> VecEsPacket;
    //    PsAacFileSource.GetNextFrame(VecEsPacket);
    //    if (VecEsPacket.empty()
    //        || VecEsPacket.begin()->BufferList.Empty()) {
    //        break;
    //    }
    //}

    PsH264AacFileSource.StopMediaSource();
    PsH265AacFileSource.StopMediaSource();
    PsH264FileSource.StopMediaSource();
    PsH265FileSource.StopMediaSource();
    //PsAacFileSource.StopMediaSource();

}

TEST(MkPsUt, PsStreamSourceTest)
{
    MkString strModuleName;
    CMkFile::GetModuleFilePath(strModuleName);
    MkString H264AacPsFileName = strModuleName + "/h264aac.ps";
    MkString H264PsFileName = strModuleName + "/h264.ps";
    MkString H265PsFileName = strModuleName + "/h265.ps";
    MkString AacPsFileName = strModuleName + "/aac.ps";
    MkString H265AacPsFileName = strModuleName + "/h265aac.ps";
    CMkPsStreamSource PsStreamSource;
    FILE *fpH264AacStream = fopen(H264AacPsFileName.c_str(), "rb");
    FILE *fpH264PsStream = fopen(H264PsFileName.c_str(), "rb");
    FILE *fpH265PsStream = fopen(H265PsFileName.c_str(), "rb");
    //FILE *fpAacPsr = fopen(AacPsFileName.c_str(), "rb");
    FILE *fpH265AacPsStream = fopen(H265AacPsFileName.c_str(), "rb");
    Uint8 pStreamBuf[LenHalfM] = { 0 };
    Uint32 nUsed = 0;
    Uint32 nOffset = 0;
    while (1) {
        int nRet = fread(pStreamBuf + nOffset, 1, LenHalfM - nOffset, fpH264AacStream);
        Uint32 nLen = nOffset + nRet;
        if (NoneError != PsStreamSource.InputData(pStreamBuf, nLen, nUsed)
            || 0 == nUsed) {
            break;
        }
        memmove(pStreamBuf, pStreamBuf + nUsed, nLen - nUsed);
        nOffset = nLen - nUsed;
    }
    EXPECT_EQ(MkCodecH264, PsStreamSource.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecAac, PsStreamSource.GetAudioParam().GetCodec());

    CMkPsStreamSource PsStreamSource1;
    nUsed = 0;
    nOffset = 0;
    while (1) {
        int nRet = fread(pStreamBuf + nOffset, 1, LenHalfM - nOffset, fpH264PsStream);
        Uint32 nLen = nOffset + nRet;
        if (NoneError != PsStreamSource1.InputData(pStreamBuf, nLen, nUsed)
            || 0 == nUsed) {
            break;
        }
        memmove(pStreamBuf, pStreamBuf + nUsed, nLen - nUsed);
        nOffset = nLen - nUsed;
    }
    EXPECT_EQ(MkCodecH264, PsStreamSource1.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecUnknown, PsStreamSource1.GetAudioParam().GetCodec());

    CMkPsStreamSource PsStreamSource2;
    nUsed = 0;
    nOffset = 0;
    while (1) {
        int nRet = fread(pStreamBuf + nOffset, 1, LenHalfM - nOffset, fpH265PsStream);
        Uint32 nLen = nOffset + nRet;
        if (NoneError != PsStreamSource2.InputData(pStreamBuf, nLen, nUsed)
            || 0 == nUsed) {
            break;
        }
        memmove(pStreamBuf, pStreamBuf + nUsed, nLen - nUsed);
        nOffset = nLen - nUsed;
    }
    EXPECT_EQ(MkCodecH265, PsStreamSource2.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecUnknown, PsStreamSource2.GetAudioParam().GetCodec());

    CMkPsStreamSource PsStreamSource3;
    nUsed = 0;
    nOffset = 0;
    while (1) {
        int nRet = fread(pStreamBuf + nOffset, 1, LenHalfM - nOffset, fpH265AacPsStream);
        Uint32 nLen = nOffset + nRet;
        if (NoneError != PsStreamSource3.InputData(pStreamBuf, nLen, nUsed)
            || 0 == nUsed) {
            break;
        }
        memmove(pStreamBuf, pStreamBuf + nUsed, nLen - nUsed);
        nOffset = nLen - nUsed;
    }
    EXPECT_EQ(MkCodecH265, PsStreamSource3.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecAac, PsStreamSource3.GetAudioParam().GetCodec());

    fclose(fpH264AacStream);
    fclose(fpH264PsStream);
    fclose(fpH265PsStream);
    fclose(fpH265AacPsStream);
}


TEST(MkPsUt, PsStreamInputSliceTest)
{
    MkString strModuleName;
    CMkFile::GetModuleFilePath(strModuleName);
    MkString H264AacPsFileName = strModuleName + "/h264aac.ps";
    MkString H264PsFileName = strModuleName + "/h264.ps";
    MkString H265PsFileName = strModuleName + "/h265.ps";
    MkString AacPsFileName = strModuleName + "/aac.ps";
    MkString H265AacPsFileName = strModuleName + "/h265aac.ps";
    CMkPsStreamSource PsStreamSource;
    FILE *fpH264AacStream = fopen(H264AacPsFileName.c_str(), "rb");
    FILE *fpH264PsStream = fopen(H264PsFileName.c_str(), "rb");
    FILE *fpH265PsStream = fopen(H265PsFileName.c_str(), "rb");
    //FILE *fpAacPsr = fopen(AacPsFileName.c_str(), "rb");
    FILE *fpH265AacPsStream = fopen(H265AacPsFileName.c_str(), "rb");
    Uint8 pStreamBuf[LenHalfM] = { 0 };
    Uint32 nUsed = 0;
    Uint32 nOffset = 0;
    while (1) {
        Uint32 nToRead = LenHalfM - nOffset;
        int nRead = fread(pStreamBuf + nOffset, 1, nToRead, fpH264AacStream);
        if (nRead <= 0) {
            break;
        }
        nOffset += nRead;
        Uint8* pNextPsHeader = nullptr;
        Uint8* pLastHeader = pStreamBuf;
        CMkPsDemuxer::GetPsHeader(pLastHeader, nOffset - (pLastHeader - pStreamBuf), &pNextPsHeader);
        pLastHeader = pNextPsHeader;
        while (NoneError == CMkPsDemuxer::GetPsHeader(pLastHeader + 4, nOffset - (pLastHeader - pStreamBuf) - 4, &pNextPsHeader)) {
            Uint8 *pOffset = pLastHeader;
            while (pOffset < pNextPsHeader) {
                Uint32 nInputLen = MkMin(pNextPsHeader - pOffset, MtuSize);
                PsStreamSource.InputDataSlice(pOffset, nInputLen, 0, nInputLen != MtuSize);
                pOffset += nInputLen;
            }
            pLastHeader = pNextPsHeader;
        }
    }
    EXPECT_EQ(MkCodecH264, PsStreamSource.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecAac, PsStreamSource.GetAudioParam().GetCodec());

    CMkPsStreamSource PsStreamSource1;
    nUsed = 0;
    nOffset = 0;
    while (1) {
        Uint32 nToRead = LenHalfM - nOffset;
        int nRead = fread(pStreamBuf + nOffset, 1, nToRead, fpH264PsStream);
        if (nRead <= 0) {
            break;
        }
        nOffset += nRead;
        Uint8* pNextPsHeader = nullptr;
        Uint8* pLastHeader = pStreamBuf;
        CMkPsDemuxer::GetPsHeader(pLastHeader, nOffset - (pLastHeader - pStreamBuf), &pNextPsHeader);
        pLastHeader = pNextPsHeader;
        while (NoneError == CMkPsDemuxer::GetPsHeader(pLastHeader + 4, nOffset - (pLastHeader - pStreamBuf) - 4, &pNextPsHeader)) {
            Uint8 *pOffset = pLastHeader;
            while (pOffset < pNextPsHeader) {
                Uint32 nInputLen = MkMin(pNextPsHeader - pOffset, MtuSize);
                PsStreamSource1.InputDataSlice(pOffset, nInputLen, 0, nInputLen != MtuSize);
                pOffset += nInputLen;
            }
            pLastHeader = pNextPsHeader;
        }
    }
    EXPECT_EQ(MkCodecH264, PsStreamSource1.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecUnknown, PsStreamSource1.GetAudioParam().GetCodec());

    CMkPsStreamSource PsStreamSource2;
    nUsed = 0;
    nOffset = 0;
    while (1) {
        Uint32 nToRead = LenHalfM - nOffset;
        int nRead = fread(pStreamBuf + nOffset, 1, nToRead, fpH265PsStream);
        if (nRead <= 0) {
            break;
        }
        nOffset += nRead;
        Uint8* pNextPsHeader = nullptr;
        Uint8* pLastHeader = pStreamBuf;
        CMkPsDemuxer::GetPsHeader(pLastHeader, nOffset - (pLastHeader - pStreamBuf), &pNextPsHeader);
        pLastHeader = pNextPsHeader;
        while (NoneError == CMkPsDemuxer::GetPsHeader(pLastHeader + 4, nOffset - (pLastHeader - pStreamBuf) - 4, &pNextPsHeader)) {
            Uint8 *pOffset = pLastHeader;
            while (pOffset < pNextPsHeader) {
                Uint32 nInputLen = MkMin(pNextPsHeader - pOffset, MtuSize);
                PsStreamSource2.InputDataSlice(pOffset, nInputLen, 0, nInputLen != MtuSize);
                pOffset += nInputLen;
            }
            pLastHeader = pNextPsHeader;
        }
    }
    EXPECT_EQ(MkCodecH265, PsStreamSource2.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecUnknown, PsStreamSource2.GetAudioParam().GetCodec());

    CMkPsStreamSource PsStreamSource3;
    nUsed = 0;
    nOffset = 0;
    while (1) {
        Uint32 nToRead = LenHalfM - nOffset;
        int nRead = fread(pStreamBuf + nOffset, 1, nToRead, fpH265AacPsStream);
        if (nRead <= 0) {
            break;
        }
        nOffset += nRead;
        Uint8* pNextPsHeader = nullptr;
        Uint8* pLastHeader = pStreamBuf;
        CMkPsDemuxer::GetPsHeader(pLastHeader, nOffset - (pLastHeader - pStreamBuf), &pNextPsHeader);
        pLastHeader = pNextPsHeader;
        while (NoneError == CMkPsDemuxer::GetPsHeader(pLastHeader + 4, nOffset - (pLastHeader - pStreamBuf) - 4, &pNextPsHeader)) {
            Uint8 *pOffset = pLastHeader;
            while (pOffset < pNextPsHeader) {
                Uint32 nInputLen = MkMin(pNextPsHeader - pOffset, MtuSize);
                PsStreamSource3.InputDataSlice(pOffset, nInputLen, 0, nInputLen != MtuSize);
                pOffset += nInputLen;
            }
            pLastHeader = pNextPsHeader;
        }
    }
    EXPECT_EQ(MkCodecH265, PsStreamSource3.GetVideoParam().GetCodec());
    EXPECT_EQ(MkCodecAac, PsStreamSource3.GetAudioParam().GetCodec());

    fclose(fpH264AacStream);
    fclose(fpH264PsStream);
    fclose(fpH265PsStream);
    fclose(fpH265AacPsStream);
}


TEST(MkPsUt, PsDeleteFileTest)
{
    MkString strModuleName;
    CMkFile::GetModuleFilePath(strModuleName);
    MkString H264AacPsFileName = strModuleName + "/h264aac.ps";
    MkString H264PsFileName = strModuleName + "/h264.ps";
    MkString H265PsFileName = strModuleName + "/h265.ps";
    MkString AacPsFileName = strModuleName + "/aac.ps";
    MkString H265AacPsFileName = strModuleName + "/h265aac.ps";
    CMkFile::DeleteLocalFile(H264AacPsFileName);
    CMkFile::DeleteLocalFile(H265AacPsFileName);
    CMkFile::DeleteLocalFile(H264PsFileName);
    CMkFile::DeleteLocalFile(H265PsFileName);
    CMkFile::DeleteLocalFile(AacPsFileName);
}